package com.lwq.predicate;

import java.nio.charset.StandardCharsets;
import java.util.HashMap;
import java.util.Iterator;
import java.util.Map;

import org.json.JSONObject;
import org.springframework.cloud.gateway.filter.GatewayFilter;
import org.springframework.cloud.gateway.filter.factory.RequestRateLimiterGatewayFilterFactory;
import org.springframework.cloud.gateway.filter.ratelimit.KeyResolver;
import org.springframework.cloud.gateway.filter.ratelimit.RateLimiter;
import org.springframework.cloud.gateway.route.Route;
import org.springframework.cloud.gateway.support.HttpStatusHolder;
import org.springframework.cloud.gateway.support.ServerWebExchangeUtils;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.stereotype.Component;


import reactor.core.publisher.Mono;

/**
 * 项目名：ifp-gateway <br>
 * 包名：com.ifp.opengate.boot.factory <br>
 * 类名：IFPRequestRateLimiterGatewayFilterFactory <br>
 * 描述：
 *
 * @author liulx
 * @create 2021-06-11 09:17
 */
@Component
public class IfpRequestRateLimiterGatewayFilterFactory extends RequestRateLimiterGatewayFilterFactory {

    private final RateLimiter defaultRateLimiter;

    private final KeyResolver defaultKeyResolver;

    private boolean denyEmptyKey = true;

    public IfpRequestRateLimiterGatewayFilterFactory(RateLimiter defaultRateLimiter, KeyResolver defaultKeyResolver) {
        super(defaultRateLimiter, defaultKeyResolver);
        this.defaultKeyResolver = defaultKeyResolver;
        this.defaultRateLimiter = defaultRateLimiter;

    }

    @Override
    public GatewayFilter apply(Config config) {
        KeyResolver resolver = (KeyResolver)this.getOrDefault(config.getKeyResolver(), this.defaultKeyResolver);
        RateLimiter<Object> limiter = (RateLimiter)this.getOrDefault(config.getRateLimiter(), this.defaultRateLimiter);
        boolean denyEmpty = (Boolean)this.getOrDefault(config.getDenyEmptyKey(), this.denyEmptyKey);
        HttpStatusHolder emptyKeyStatus =
            HttpStatusHolder.parse((String)this.getOrDefault(config.getEmptyKeyStatus(), this.getEmptyKeyStatusCode()));
        return (exchange, chain) -> {
            Route route = (Route)exchange.getAttribute(ServerWebExchangeUtils.GATEWAY_ROUTE_ATTR);
            return resolver.resolve(exchange).defaultIfEmpty("____EMPTY_KEY__").flatMap((key) -> {
                if ("____EMPTY_KEY__".equals(key)) {
                    if (denyEmpty) {
                        ServerWebExchangeUtils.setResponseStatus(exchange, emptyKeyStatus);
                        return exchange.getResponse().setComplete();
                    } else {
                        return chain.filter(exchange);
                    }
                } else {
                    return limiter.isAllowed(route.getId(), key).flatMap((response) -> {
                        Iterator var4 = response.getHeaders().entrySet().iterator();

                        while (var4.hasNext()) {
                            Map.Entry<String, String> header = (Map.Entry)var4.next();
                            exchange.getResponse().getHeaders().add((String)header.getKey(), (String)header.getValue());
                        }

                        if (response.isAllowed()) {
                            return chain.filter(exchange);
                        } else {
                            // ServerWebExchangeUtils.setResponseStatus(exchange, config.getStatusCode());
                            // return exchange.getResponse().setComplete();
                            System.out.println("服务已限流");
                            /*return reactor.core.publisher.Mono.error(new AccessControlException(
                                    ErrorEnum.merchantCurrentCotrolError.code(), ErrorEnum.merchantCurrentCotrolError.msg()
                            ));*/
                            // TODO: 在这里添加触发限流后通知的警告！
                            ServerHttpResponse httpResponse = exchange.getResponse();
                            httpResponse.setStatusCode(config.getStatusCode());

                            Map<String, Object> dataMap = new HashMap<>(4);
                            dataMap.put("errorCode", 10000001);
                            dataMap.put("errorMsg", "服务已限流");

                            // DataBuffer buffer =
                            // httpResponse.bufferFactory().wrap(data.getBytes(StandardCharsets.UTF_8));
                            DataBuffer buffer = httpResponse.bufferFactory()
                                .wrap(JSONObject.wrap(dataMap).toString().getBytes(StandardCharsets.UTF_8));
                            return httpResponse.writeWith(Mono.just(buffer));
                        }
                    });
                }
            });
        };
    }

    private <T> T getOrDefault(T configValue, T defaultValue) {
        return configValue != null ? configValue : defaultValue;
    }

}
